#!/usr/bin/env python3
# -*- coding: utf8 -*-

import json
import os
import re

OUTFILE = 'out.s'
PROGRAM_ERROR_OUTPUT = 'err.out'
RARS_OUTPUT = "rars.out"
RARS_EXITCODE = "rars.exitcode"
RARS_TIME_LIMIT = 60
RARS = "rars.jar"

TEST_DIR = './tests'
EXP_CASES_DIR = TEST_DIR + '/exp'
VAR_CASES_DIR = TEST_DIR + '/var'
REG_CASES_DIR = TEST_DIR + '/reg'
CONTROLFLOW_CASES_DIR = TEST_DIR + '/controlflow'
RESULT = {"score": 0, "details": []}


def extract_instructions(text):
    # 匹配所有指令执行的数量
    all_executed = re.search(r"All instructions executed: (\d+)", text)
    if all_executed:
        all_executed = int(all_executed.group(1))
    else:
        all_executed = None

    # 匹配Load指令执行的数量
    load_executed = re.search(r"Load instructions executed: (\d+)", text)
    if load_executed:
        load_executed = int(load_executed.group(1))
    else:
        load_executed = None

    # 匹配Store指令执行的数量
    store_executed = re.search(r"Store instructions executed: (\d+)", text)
    if store_executed:
        store_executed = int(store_executed.group(1))
    else:
        store_executed = None

    return all_executed, load_executed, store_executed


def check_asm_answer(expect_result, check_reg=False, limit=0):
    cmd = f' gtimeout {RARS_TIME_LIMIT} java -jar {RARS} {OUTFILE} ic >{RARS_OUTPUT} 2>&1; echo $? > {RARS_EXITCODE}'

    cmd_result = os.system(cmd)
    output_msg = open(RARS_OUTPUT).read()
    if cmd_result == 124:
        score = 0
        message = "timeout"
    else:
        your_answer = int(open(RARS_EXITCODE).read().strip())
        expect_result_short = ((expect_result % 256) + 256) % 256
        print(your_answer, expect_result, expect_result_short)
        if your_answer != expect_result and your_answer != expect_result_short:
            message = "wrong answer."
            score = 0
        else:
            if not check_reg:
                message = "correct!"
                score = 100
            else:
                all_executed, load_executed, store_executed = extract_instructions(output_msg)
                if load_executed + store_executed > limit:
                    total = load_executed + store_executed
                    times = total / limit
                    message = f"Your answer is correct, but load/store count is {times:.2f} times the limit."
                    score = min(100, int(100 * (1.0 / times)))  # Score between 0-100
                    if score <= 20:
                        score = 0
                else:
                    score = 100
                    message = "correct!"
    return score, message


def check_dir(dir_name):
    files = os.listdir(dir_name)
    files = [file_name for file_name in files if file_name.endswith('.sy')]
    files.sort()
    
    part_score = 0
    first_error_msg = "All tests passed"
    
    for file_name in files:
        answer_file_name = file_name.replace('.sy', '.out')
        expect_result = int(open(dir_name + "/" + answer_file_name).read().strip())
        os.remove(dir_name + "/" + answer_file_name)
        
        limit = None
        if 'reg' in dir_name or 'controlflow' in dir_name:
            limit_file_name = file_name.replace('.sy', '.limit')
            limit = int(open(dir_name + "/" + limit_file_name).read().strip())
            os.remove(dir_name + "/" + limit_file_name)

        cmd = f'../target/release/compiler {dir_name + "/" + file_name} {OUTFILE} 1> /dev/null 2> {PROGRAM_ERROR_OUTPUT}'
        cmd_result = os.system(cmd)
        program_stderr_str = open(PROGRAM_ERROR_OUTPUT).read()
        
        if cmd_result != 0:
            score = 0
            message = program_stderr_str.replace('\n', ' ', 2).split('\n')[0].replace('\t', ' ')
            message = message.split(':')[0]
            if message.startswith('Exception in thread "main" ') and message.endswith('Exception'):
                message = 'Exception in thread "main"'
            else:
                message = f'Runtime Error: {program_stderr_str}'
            
            if first_error_msg == "All tests passed":
                first_error_msg = f"{message}"
            part_score += score
            continue
            
        elif program_stderr_str.strip() != '':
            score = 0
            message = 'stderr output detected'
            
            if first_error_msg == "All tests passed":
                first_error_msg = f"{message}"
            part_score += score
            continue
            
        elif open(OUTFILE).read().strip() == '':
            score = 0
            message = 'Empty risc-v asm output.'
            
            if first_error_msg == "All tests passed":
                first_error_msg = f"{message}"

            part_score += score
            continue
            
        else:
            if 'reg' in dir_name or 'controlflow' in dir_name:
                score, message = check_asm_answer(expect_result, True, limit)
            else:
                score, message = check_asm_answer(expect_result)
            
            if score < 100 and first_error_msg == "All tests passed":
                first_error_msg = f"{message}"
            if message == "wrong answer." and first_error_msg != "wrong answer.":
                first_error_msg = f"{message}"
            
            part_score += score
    
    total_score = len(files) * 100
    
    RESULT['score'] += part_score
    RESULT['details'].append(
            {"score": part_score, "message":  f"{part_score}/{total_score}, {first_error_msg}", "title": os.path.basename(dir_name)})


def main():
    check_dir('./tests/exp')
    check_dir('./tests/var')
    check_dir('./tests/reg')
    check_dir('./tests/controlflow')
    print(json.dumps(RESULT))


if __name__ == '__main__':
    try:
        main()
    except BaseException as e:
        print(json.dumps({"score": 0, "details": [
            {"score": 0, "title": f"Exception in OJ scripts {e}", "message": "Please contact your TA"}]}))
